"""
connection.py
=============
Neo4j Database Connection Utilities

Provides:
- Connection pooling
- Session management  
- Retry logic
- Health checks

Author: Neo4j Tutorial
Date: January 2025
"""

import os
from contextlib import contextmanager
from typing import Optional, Dict, Any, List, Generator
from dotenv import load_dotenv
from neo4j import GraphDatabase, Driver, Session, Result
from neo4j.exceptions import ServiceUnavailable, AuthError, TransientError
import time

# Load environment variables
load_dotenv()


class Neo4jConnection:
    """
    A robust Neo4j connection manager with connection pooling
    and automatic retry logic.
    """
    
    def __init__(
        self,
        uri: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        database: str = "neo4j",
        max_connection_pool_size: int = 50,
        connection_timeout: float = 30.0,
        max_retry_time: float = 30.0
    ):
        """
        Initialize Neo4j connection.
        
        Args:
            uri: Neo4j URI (bolt://localhost:7687)
            user: Database username
            password: Database password
            database: Database name (default: neo4j)
            max_connection_pool_size: Maximum connections in pool
            connection_timeout: Connection timeout in seconds
            max_retry_time: Maximum retry time for transient errors
        """
        self.uri = uri or os.getenv("NEO4J_URI", "bolt://localhost:7687")
        self.user = user or os.getenv("NEO4J_USER", "neo4j")
        self.password = password or os.getenv("NEO4J_PASSWORD", "password")
        self.database = database or os.getenv("NEO4J_DATABASE", "neo4j")
        
        self._driver: Optional[Driver] = None
        self._max_pool_size = max_connection_pool_size
        self._connection_timeout = connection_timeout
        self._max_retry_time = max_retry_time
        
    def connect(self) -> Driver:
        """
        Establish connection to Neo4j database.
        
        Returns:
            Neo4j Driver instance
            
        Raises:
            AuthError: If authentication fails
            ServiceUnavailable: If database is unreachable
        """
        if self._driver is None:
            try:
                self._driver = GraphDatabase.driver(
                    self.uri,
                    auth=(self.user, self.password),
                    max_connection_pool_size=self._max_pool_size,
                    connection_timeout=self._connection_timeout,
                    max_transaction_retry_time=self._max_retry_time
                )
                # Verify connectivity
                self._driver.verify_connectivity()
                print(f"✅ Connected to Neo4j at {self.uri}")
                
            except AuthError as e:
                print(f"❌ Authentication failed: {e}")
                raise
            except ServiceUnavailable as e:
                print(f"❌ Database unavailable: {e}")
                raise
                
        return self._driver
    
    def close(self):
        """Close the driver connection."""
        if self._driver:
            self._driver.close()
            self._driver = None
            print("🔌 Connection closed")
    
    @contextmanager
    def session(self, database: Optional[str] = None) -> Generator[Session, None, None]:
        """
        Context manager for Neo4j sessions.
        
        Args:
            database: Override default database
            
        Yields:
            Neo4j Session
        """
        driver = self.connect()
        session = driver.session(database=database or self.database)
        try:
            yield session
        finally:
            session.close()
    
    def execute_query(
        self,
        query: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute a Cypher query and return results.
        
        Args:
            query: Cypher query string
            parameters: Query parameters
            database: Override default database
            
        Returns:
            List of result records as dictionaries
        """
        with self.session(database) as session:
            result = session.run(query, parameters or {})
            return [record.data() for record in result]
    
    def execute_write(
        self,
        query: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Execute a write query within a transaction.
        
        Args:
            query: Cypher write query
            parameters: Query parameters
            database: Override default database
            
        Returns:
            Query counters and summary
        """
        def _write_tx(tx, query, params):
            result = tx.run(query, params)
            summary = result.consume()
            return {
                "nodes_created": summary.counters.nodes_created,
                "nodes_deleted": summary.counters.nodes_deleted,
                "relationships_created": summary.counters.relationships_created,
                "relationships_deleted": summary.counters.relationships_deleted,
                "properties_set": summary.counters.properties_set,
            }
        
        with self.session(database) as session:
            return session.execute_write(_write_tx, query, parameters or {})
    
    def health_check(self) -> Dict[str, Any]:
        """
        Perform a health check on the database.
        
        Returns:
            Dictionary with health status and database info
        """
        try:
            with self.session() as session:
                # Get database info
                result = session.run("""
                    CALL dbms.components() YIELD name, versions, edition
                    RETURN name, versions, edition
                """)
                record = result.single()
                
                # Get node and relationship counts
                counts = session.run("""
                    MATCH (n) WITH count(n) as nodeCount
                    OPTIONAL MATCH ()-[r]->()
                    RETURN nodeCount, count(r) as relCount
                """).single()
                
                return {
                    "status": "healthy",
                    "name": record["name"],
                    "version": record["versions"][0],
                    "edition": record["edition"],
                    "node_count": counts["nodeCount"],
                    "relationship_count": counts["relCount"]
                }
                
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }
    
    def get_database_stats(self) -> Dict[str, Any]:
        """
        Get detailed database statistics.
        
        Returns:
            Database statistics
        """
        with self.session() as session:
            # Get label counts
            labels_result = session.run("""
                CALL db.labels() YIELD label
                CALL {
                    WITH label
                    MATCH (n) WHERE label IN labels(n)
                    RETURN count(n) as count
                }
                RETURN label, count
                ORDER BY count DESC
            """)
            labels = {r["label"]: r["count"] for r in labels_result}
            
            # Get relationship type counts
            rel_result = session.run("""
                CALL db.relationshipTypes() YIELD relationshipType
                CALL {
                    WITH relationshipType
                    MATCH ()-[r]->()
                    WHERE type(r) = relationshipType
                    RETURN count(r) as count
                }
                RETURN relationshipType, count
                ORDER BY count DESC
            """)
            relationships = {r["relationshipType"]: r["count"] for r in rel_result}
            
            # Get index count
            index_result = session.run("SHOW INDEXES YIELD name RETURN count(*) as count")
            index_count = index_result.single()["count"]
            
            return {
                "labels": labels,
                "relationships": relationships,
                "index_count": index_count,
                "total_nodes": sum(labels.values()) if labels else 0,
                "total_relationships": sum(relationships.values()) if relationships else 0
            }
    
    def __enter__(self):
        """Context manager entry."""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()


# Singleton connection instance
_connection: Optional[Neo4jConnection] = None


def get_connection() -> Neo4jConnection:
    """Get or create the singleton connection instance."""
    global _connection
    if _connection is None:
        _connection = Neo4jConnection()
    return _connection


def get_driver() -> Driver:
    """Get the Neo4j driver instance."""
    return get_connection().connect()


def reset_connection():
    """Reset the singleton connection (useful for testing)."""
    global _connection
    if _connection:
        _connection.close()
    _connection = None


# ============================================================
# TESTING
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Testing Neo4j Connection")
    print("=" * 60)
    
    # Test connection
    with Neo4jConnection() as conn:
        health = conn.health_check()
        print(f"\nHealth Status: {health['status']}")
        
        if health['status'] == 'healthy':
            print(f"Neo4j Version: {health['version']}")
            print(f"Edition: {health['edition']}")
            print(f"Nodes: {health['node_count']}")
            print(f"Relationships: {health['relationship_count']}")
            
            # Test query
            result = conn.execute_query("RETURN 'Connection successful!' as message")
            print(f"\nTest Query: {result[0]['message']}")
    
    print("\n✅ Connection test complete!")
